import pandas as pd
from dataclasses import dataclass
from typing import Dict, Tuple
from functools import cached_property
from zkl_aiutils_datasets import load_dataset
from zkl_pyutils_fsspec import resolve_fs

# @dataclass
# class NhanesValueMap:
#     value_map: Dict[str, int]

@dataclass(kw_only=True)
class NhanesEncoding:
    new_code: int
    # 此处为new_code对应的意思，表格里是label列
    meaning: str

# @dataclass(kw_only=True)
# class UKBioBankField:
#     columns_name: str
#     description: str


class NhanesEncodingProcessor:
    def __init__(self, codebook_path: str = None, dictionary_path: str = None, dataset_path: str = None):
        self.codebook_path = codebook_path
        self.dictionary_path = dictionary_path
        self.dataset_path = dataset_path

        if codebook_path:
            self.codebook = pd.read_excel(codebook_path)
        if dictionary_path:
            self.dictionary = pd.read_excel(dictionary_path)
        
    @cached_property
    def value_maps(self) -> Dict[Tuple[str, str], Dict[str, int]]:
        value_maps: Dict[Tuple[str, str], Dict[str, int]] = {}
        for _, row in self.codebook.iterrows():
            item = str(row['Item']).lower()
            variable = str(row['variable'])
            code = str(row['code'])
            new_code = int(row['new_code'])
            key = (item, variable)
            if key not in value_maps:
                value_maps[key] = {}
            value_maps[key][code] = new_code
        return value_maps

    @cached_property
    def new_code_to_label(self) -> Dict[int, str]:
        new_code_to_label: Dict[int, str] = {}
        for _, row in self.codebook.iterrows():
            new_code = int(row['new_code'])
            label = str(row['label'])
            if new_code not in new_code_to_label:
                new_code_to_label[new_code] = label
        return new_code_to_label

    def get_value_map(self, item: str, variable: str) -> Dict[str, int]:
        return self.value_maps.get((item.lower(), variable), {})


    def is_categorical(self, item: str, variable: str) -> bool:
        return bool(self.get_value_map(item, variable))

    # 建立从new_code到label的映射，仿造omics_bert_label_encoding/utils/processing_uk.py
    @cached_property
    def encodings(self) -> dict[int, NhanesEncoding]:
        encodings_map = {}
        for new_code, label in self.new_code_to_label.items():
            encodings_map[new_code] = NhanesEncoding(new_code=new_code, meaning=label)
        return encodings_map

    # 建立从columns_name到description的映射
    @cached_property
    def columns_name_to_description(self) -> dict[str, str]:
        columns_name_to_description = {}
        # self.dictionary 是一个excel表格，有两列，分别是variable_codename_use和variable_description_use，对应columns_name和description
        # self.dictionary_path去重，仅保留第一次出现的column_name和对应的description
        self.dictionary = self.dictionary.drop_duplicates(subset=['variable_codename_use'])
        for _, row in self.dictionary.iterrows():
            column_name = str(row['variable_codename_use'])
            description = str(row['variable_description_use'])
            columns_name_to_description[column_name] = description
        return columns_name_to_description

    @cached_property
    def get_column_data_types(self) -> Dict[str, str]:
        if not self.dataset_path:
            return {}
        fs = resolve_fs(self.dataset_path)
        dataset = load_dataset(fs)
        types = {}
        for type_name, child in dataset.named_children.items():
            for col in child.columns_name:
                if type_name == "str_columns":
                    types[col] = "string"
                elif type_name == "float_columns":
                    types[col] = "float"
                elif type_name == "integer_columns":
                    if col == "eid":
                        types[col] = "integer"
                    else:
                        types[col] = "category"
        return types

if __name__ == "__main__":
    # 测试样例
    dataset_path = "/home/xw/python3test/omics-bert-label-encoding/datasets/ready/nhanes/v1"
    test_processor = NhanesEncodingProcessor(dataset_path=dataset_path)
    column_types = test_processor.get_column_data_types
    print("列数据类型映射:", column_types)

    
